# -*- coding: utf-8 -*-
import random
import numpy as np
import ConfigSpace as CS
import ConfigSpace.hyperparameters as CSH
from onnx import TensorProto
from onnx import helper
from .Conv import Conv

class QLinearConv(Conv):
    def __init__(self):
        print('QLinearConv Init--------------------------')
        super().__init__()

    def get_initializer_name_list(self):
            initializer_name_list = ['x_scale', 'x_zero_point', 'w_scale', 'w_zero_point', 'y_scale', 'y_zero_point']
            return initializer_name_list

    def generate_random_initializer(self, in_shapes, out_shape, dtype):
        zp_type = TensorProto.UINT8
        if dtype == 'tensor(int8)': 
            zp_type = TensorProto.INT8

        x_scale_shape = [1]
        x_zp_shape = [1]
        v1 = np.random.randint(0, 100)
        x_scale_array = np.random.random_sample(1)
        #x_scale_array = v1*x_scale_array
        x_zp_array = [np.random.randint(0, 256)]
        if dtype == 'tensor(int8)':
            x_zp_array = [np.random.randint(-127, 128+1)]

        w_scale_shape = [1]
        w_zp_shape = [1]
        v2 = np.random.randint(0, 100)
        w_scale_array = np.random.random_sample(1)
        #w_scale_array = v2*w_scale_array
        w_zp_array = [np.random.randint(0, 256)]
        if dtype == 'tensor(int8)':
            w_zp_array = [np.random.randint(-127, 128+1)]
    
        y_scale_shape = [1]
        y_zp_shape = [1]
        v3 = np.random.randint(0, 100)
        y_scale_array = np.random.random_sample(1)
        y_scale_array = v3*y_scale_array
        y_zp_array = [np.random.randint(0, 256)]
        if dtype == 'tensor(int8)':
            y_zp_array = [np.random.randint(-127, 128+1)]

        '''
        r = random.randint(1, 5)
        if r % 2 == 0:
            x_scale_shape = in_shapes[0].copy()
            x_zp_shape = in_shapes[0].copy()

            size = 1
            for d in x_scale_shape:
                size = size * d
            x_scale_array = np.random.random_sample(size)
            #x_scale_array = float(v1*x_scale_array)

            x_zp_array = np.random.randint(0, 256, size)
            if dtype == 'tensor(int8)':
                x_zp_array = np.random.randint(-127, 128+1, size)

            w_scale_shape = in_shapes[1].copy()
            w_zp_shape = in_shapes[1].copy()

            size = 1
            for d in w_scale_shape:
                size = size * d
            w_scale_array = np.random.random_sample(size)
            #w_scale_array = v2*w_scale_array
            w_zp_array = np.random.randint(0, 256, size)
            if dtype == 'tensor(int8)':
                w_zp_array = np.random.randint(-127, 128+1, size)

            y_scale_shape = out_shape.copy()
            y_zp_shape = out_shape.copy()

            size = 1
            for d in y_scale_shape:
                size = size * d

            y_scale_array = np.random.random_sample(size)
            #y_scale_array = v3*y_scale_array
            y_zp_array = np.random.randint(0, 256, size)
            if dtype == 'tensor(int8)':
                y_zp_array = np.random.randint(-127, 128+1, size)    
        '''

        print('x_scale_shape: ', x_scale_shape)
        print('x_scale_array: ', x_scale_array)
        print('x_zp_shape: ', x_zp_shape)
        print('x_zp_array: ', x_zp_array)

        print('w_scale_shape: ', w_scale_shape)
        print('w_scale_array: ', w_scale_array)
        print('w_zp_shape: ', w_zp_shape)
        print('w_zp_array: ', w_zp_array)

        print('y_scale_shape: ', y_scale_shape)
        print('y_scale_array: ', w_scale_array)
        print('y_zp_shape: ', y_zp_shape)
        print('y_zp_array: ', y_zp_array)

        x_scale = helper.make_tensor('x_scale', TensorProto.FLOAT, x_scale_shape, x_scale_array)    
        x_zero_point = helper.make_tensor('x_zero_point', zp_type, x_zp_shape, x_zp_array) 

        w_scale = helper.make_tensor('w_scale', TensorProto.FLOAT, w_scale_shape, w_scale_array)    
        w_zero_point = helper.make_tensor('w_zero_point', zp_type, w_zp_shape, w_zp_array)

        y_scale = helper.make_tensor('y_scale', TensorProto.FLOAT, y_scale_shape, y_scale_array)    
        y_zero_point = helper.make_tensor('y_zero_point', zp_type, y_zp_shape, y_zp_array)

        initializer = [x_scale, x_zero_point, w_scale, w_zero_point, y_scale, y_zero_point]

        return initializer

    def get_input_name(self):
        input_name = ['x', 'x_scale', 'x_zero_point', 'w', 'w_scale', 'w_zero_point', 'y_scale', 'y_zero_point', 'B']
        return input_name

    def get_input_type(self, name, type):
        if name == 'B':
            return 'tensor(int32)'

        return type    



